import scipy.io as sio
import numpy as np
from time import strftime
import config_ql as qc
import os
from helpers.logger_maker import get_logger
import matplotlib.pyplot as plt

def savitzky_golay(y, window_size, order, deriv=0, rate=1):
    from math import factorial
    try:
        window_size = np.abs(np.int(window_size))
        order = np.abs(np.int(order))
    except ValueError:
        raise ValueError("window_size and order have to be of type int")
    if window_size % 2 != 1 or window_size < 1:
        raise TypeError("window_size size must be a positive odd number")
    if window_size < order + 2:
        raise TypeError("window_size is too small for the polynomials order")
    order_range = range(order + 1)
    half_window = (window_size - 1) // 2
    # precompute coefficients
    b = np.mat([[k ** i for i in order_range] for k in range(-half_window, half_window + 1)])
    m = np.linalg.pinv(b).A[deriv] * rate ** deriv * factorial(deriv)
    firstvals = y[0] - np.abs(y[1:half_window + 1][::-1] - y[0])
    lastvals = y[-1] + np.abs(y[-half_window - 1:-1][::-1] - y[-1])
    y = np.concatenate((firstvals, y, lastvals))
    return np.convolve(m[::-1], y, mode='valid')

class DataSaver:
    def __init__(self, saver_name):
        #self.logger = get_logger("DataSaver_%s" % (saver_name))

        self.saver_name = saver_name
        self.time_str = strftime("%Y%m%d_%H%M%S")
        self.dir_name = "%s_%s" % (self.saver_name, self.time_str)
        self.data_dir = os.path.join("results", self.dir_name)
        if not os.path.exists(self.data_dir):
            os.mkdir(self.data_dir)

        self.save_RMSE = np.zeros((qc.EPISODE_NUMBER, qc.STEP_NUMBER))
        self.save_Energy_Consumption = np.zeros((qc.EPISODE_NUMBER, qc.STEP_NUMBER))
        self.save_Utility = np.zeros((qc.EPISODE_NUMBER, qc.STEP_NUMBER))

        self.save_dict = {
            "RMSE": self.save_RMSE,
            "Energy_Consumption": self.save_Energy_Consumption,
            "Utility": self.save_Utility,
        }

    def save_step(self, episode, step, Utility, state_results):
        RMSE = state_results['Localization_Error']
        Energy_Consumption = state_results['beacons'] * 2

        self.save_RMSE[episode, step] = RMSE
        self.save_Energy_Consumption[episode, step] = Energy_Consumption
        self.save_Utility[episode, step] = Utility

        '''
        self.logger.info("EP:%-3d ST:%-4d JC:%02d PO:%-3d LO:%-3d DL:%-8.4f WL:%-8.4f CH:%02d",
                    episode, step, jamming_channel, position, packet_loss, data_latency, wifi_latency, channel)
        '''

    def save_qtable(self, episode, ql_agent):
        qtable_dir = os.path.join(self.data_dir, "Q")
        if not os.path.exists(qtable_dir):
            os.mkdir(qtable_dir)
        qtable_path = os.path.join(qtable_dir, "%d.mat" % (episode))
        ql_agent.ql.save_q(qtable_path)
        #self.logger.info("Save Q: %s", qtable_path)

    def save_mat(self):
        save_mat_path = os.path.join(self.data_dir, "results.mat")
        sio.savemat(save_mat_path, self.save_dict)
        #self.logger.info(self.dir_name)

    def plot_mat(self):
        plt.figure()
        plt.plot(savitzky_golay(self.save_RMSE.mean(axis=0), 11, 1))
        plt.title("RMSE")
        plt.figure()
        plt.plot(savitzky_golay(self.save_Energy_Consumption.mean(axis=0), 11, 1))
        plt.title("Energy_Consumption")
        plt.figure()
        plt.plot(savitzky_golay(self.save_Utility.mean(axis=0), 11, 1))
        plt.title("Utility")
        plt.show()
